from openai import OpenAI
from ..utils import *
# Set OpenAI's API key and API base to use vLLM's API server.

# export NVIDIA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
# 启动服务的命令： vllm serve /home/A02_tmpdata3/ckpt/transformers/qwen3_8B  --tensor-parallel-size 8

# openai_api_key = "EMPTY"
# openai_api_base = "http://10.21.4.2:8000/v1"

client = None

def init_vllm_qwen3_client(openai_api_base = "http://10.21.4.4:8000/v1", openai_api_key = "EMPTY"):
    print_info("Initializing vLLM's Qwen3 client...， API base: {}, API key: {}".format(openai_api_base, openai_api_key))
    global client
    client = OpenAI(
        api_key=openai_api_key,
        base_url=openai_api_base,
    )

def qwen3_vllm_chat(question, think=False):
    if client is None:
        # 将构建的base_url传入初始化函数
        init_vllm_qwen3_client()
    chat_response = client.chat.completions.create(
        model="/home/A02_tmpdata3/ckpt/transformers/qwen3_8B",
        messages=[
            {"role": "user", "content": question},
        ],
        max_tokens=8192,
        temperature=1.2,
        top_p=0.95,
        presence_penalty=1.5,
        extra_body={
            "top_k": 20,
            "chat_template_kwargs": {"enable_thinking": think},
        },
    )
    txt  = chat_response.choices[0].message.content
    return txt